package io.yanglong.aiassistant.config;

import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.router.DefaultQueryRouter;
import dev.langchain4j.rag.query.router.QueryRouter;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;
import io.yanglong.aiassistant.service.DocSearcher;
import io.yanglong.aiassistant.tools.DocTools;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration(proxyBeanMethods = false)
public class ModelConfig {

    @Bean
    public EmbeddingModel qwenEmbeddingModel() {
        EmbeddingModel embeddingModel = OllamaEmbeddingModel.builder()
                .baseUrl("http://127.0.0.1:11434")
                .modelName(Constant.MODEL_NAME_QW2_5_7B)
                .build();
        return embeddingModel;
    }

    @Bean
    public EmbeddingStore<TextSegment> redisEmbeddingStore() {
        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                //维度，根据模型设置，ollama show [modelName]
                .dimension(Constant.DS_7B_QWEN_DIMENSION)
                .indexName("doc-rag")
                .build();
        return embeddingStore;
    }

    @Bean
    public ContentRetriever qwenContentRetriever(EmbeddingStore<TextSegment> redisEmbeddingStore,
                                                 EmbeddingModel qwenEmbeddingModel) {
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(qwenEmbeddingModel)
                .embeddingStore(redisEmbeddingStore)
                //最多返回5条知识条目
                .maxResults(5)
                //返回相关评分高于此值的知识条目，只支持0.0到1.0
                .minScore(0.6d)
                .build();
        return contentRetriever;
    }

    @Bean
    public ChatLanguageModel qwenChatLanguageModel() {
        ChatLanguageModel chatLanguageModel = OllamaChatModel.builder()
                .baseUrl("http://127.0.0.1:11434")
                .modelName(Constant.MODEL_NAME_QW2_5_7B)
                .maxRetries(3)
                .build();
        return chatLanguageModel;
    }

    @Bean
    public RetrievalAugmentor retrievalAugmentor(ContentRetriever qwenContentRetriever) {
        QueryRouter queryRouter = new DefaultQueryRouter(qwenContentRetriever);
        RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                .queryRouter(queryRouter)
                .contentInjector(DefaultContentInjector.builder()
                        .promptTemplate(Constant.QA_PROMPT_TEMPLATE)
                        .build())
                .build();
        return retrievalAugmentor;
    }

    @Bean
    public DocSearcher qwenDocSearcher(ChatLanguageModel qwenChatLanguageModel, RetrievalAugmentor retrievalAugmentor) {
        DocSearcher searcher = AiServices.builder(DocSearcher.class)
                .chatLanguageModel(qwenChatLanguageModel)
                .chatMemoryProvider(id -> MessageWindowChatMemory.withMaxMessages(10))
                .tools(ToolSpecifications.toolSpecificationsFrom(DocTools.class))
                .retrievalAugmentor(retrievalAugmentor)
                .build();
        return searcher;
    }
}
